from tensorflow import keras
from sklearn.datasets import load_svmlight_file

num_classes = 10


def load_to_numpy(path):
    data = load_svmlight_file(path)
    return (data[0].toarray(), data[1])


path = "usps"
path_t = "usps.t"

(x_train, y_train) = load_to_numpy(path)
(x_test, y_test) = load_to_numpy(path_t)

x_train = x_train.reshape(x_train.shape[0], 16, 16, 1)
x_test = x_test.reshape(x_test.shape[0], 16, 16, 1)

y_train = y_train - 1
y_test = y_test - 1

y_train = keras.utils.to_categorical(y_train, num_classes)
y_test = keras.utils.to_categorical(y_test, num_classes)

# LeNet-5 model
model = keras.Sequential()
model.add(keras.layers.Conv2D(32, kernel_size=(3,3), activation='relu'))
model.add(keras.layers.MaxPooling2D(pool_size=(2,2)))
model.add(keras.layers.Conv2D(64, kernel_size=(3,3), activation='relu'))
model.add(keras.layers.MaxPooling2D(pool_size=(2,2)))
model.add(keras.layers.Flatten())
model.add(keras.layers.Dropout(0.5))
model.add(keras.layers.Dense(num_classes, activation='softmax'))


# 设置优化器、损失函数、准确率
model.compile(optimizer='adam', loss='categorical_crossentropy', metrics=['acc'])
# 载入训练集、测试级，设置训练轮次
history = model.fit(x_train, y_train, batch_size=256, epochs=30, validation_data=(x_test, y_test), validation_split = 0.15)
# 使用测试集进行评估
model.evaluate(x_test, y_test)
# 保存模型
model.save('usps.h5')
